# math_utils.py

import secrets  # 修改：使用 secrets 模块替代 random
import math
import hashlib
import os

# ... (extended_gcd, mod_inverse, power 函数保持不变) ...
def extended_gcd(a, b):
    """
    扩展欧几里得算法,迭代实现
    返回 (gcd, x, y) 使得 ax + by = gcd.
    保证返回的 gcd 是非负的。
    """
    if not isinstance(a, int) or not isinstance(b, int):
        raise TypeError("输入必须是整数")
    r_prev, r_curr = a, b
    s_prev, s_curr = 1, 0
    t_prev, t_curr = 0, 1
    while r_curr != 0:
        quotient = r_prev // r_curr
        r_prev, r_curr = r_curr, r_prev - quotient * r_curr
        s_prev, s_curr = s_curr, s_prev - quotient * s_curr
        t_prev, t_curr = t_curr, t_prev - quotient * t_curr
    gcd, x, y = r_prev, s_prev, t_prev
    if gcd < 0:
        gcd, x, y = -gcd, -x, -y
    return gcd, x, y

def mod_inverse(a, m):
    """
    计算 a 在模 m 下的乘法逆元
    即找到 x 使得 (a * x) % m == 1
    """
    gcd, x, y = extended_gcd(a, m)
    if gcd != 1:
        raise ValueError(f"{a} 和 {m} 不是互素关系，无法计算模逆元")
    return (x % m + m) % m

def power(base, exp, mod):
    """
    模幂运算 (base^exp) % mod
    使用快速幂算法
    """
    res = 1
    base %= mod
    while exp > 0:
        if exp % 2 == 1:
            res = (res * base) % mod
        base = (base * base) % mod
        exp //= 2
    return res


def is_prime_miller_rabin(n, k=40):
    """
    Miller-Rabin 素性测试
    """
    if n <= 1: return False
    if n <= 3: return True
    if n % 2 == 0: return False

    s, d = 0, n - 1
    while d % 2 == 0:
        d //= 2
        s += 1

    for _ in range(k):
        # 修改：使用 secrets 模块生成安全的随机见证数 a
        # secrets.randbelow(n-3) 生成 [0, n-4] 范围的整数, +2 后范围为 [2, n-2]
        # 这等价于 random.randrange(2, n-1)
        a = secrets.randbelow(n - 3) + 2
        x = power(a, d, n)

        if x == 1 or x == n - 1:
            continue

        for _ in range(s - 1):
            x = power(x, 2, n)
            if x == n - 1:
                break
        else:
            return False
    return True

def generate_random_bits_number(length):
    """
    生成一个指定比特长度的随机数。
    """
    if length <= 0:
        raise ValueError("比特长度必须为正。")
    # 修改：使用 secrets.randbits
    num = secrets.randbits(length - 1)
    num |= (1 << (length - 1))
    return num

def generate_prime_candidate(length):
    """
    生成一个指定位长的奇数候选者
    """
    # 修改：使用 secrets.randbits
    p = secrets.randbits(length)
    p |= (1 << (length - 1)) | 1
    return p

# ... (generate_large_prime 和 generate_prime_number 函数无需修改，它们调用的是已修改的函数) ...
def generate_large_prime(length=256):
    count = 0
    while True:
        p = generate_prime_candidate(length)
        count += 1
        if is_prime_miller_rabin(p, k=20):
            return p

def generate_prime_number(length=1024):
    if length < 8:
        raise ValueError("素数位数至少应为8位")
    p = generate_prime_candidate(length)
    while not is_prime_miller_rabin(p, 40):
        p = generate_prime_candidate(length)
    return p

# ... (generate_safe_prime 函数无需修改) ...
def generate_safe_prime(length):
    assert length >= 8
    q_length = length - 1
    attempts = 0
    while True:
        q = generate_large_prime(q_length)
        p = 2 * q + 1
        if is_prime_miller_rabin(p, 40):
            print(f"找到安全素数 p: {p} (q: {q})，共尝试了 {attempts} 次.")
            return p, q
        attempts += 1
        if attempts % 10 == 0:
            print(f"  尝试了 {attempts} 次仍未成功...")


# ... (bytes_to_int, int_to_bytes, mgf1, bytes_xor 函数保持不变) ...
def bytes_to_int(byte_data):
    return int.from_bytes(byte_data, 'big')

def int_to_bytes(number, min_length=0):
    length = (number.bit_length() + 7) // 8
    if min_length > length:
        length = min_length
    return number.to_bytes(length, 'big')

def mgf1(seed_bytes, mask_len_bytes, hash_func=hashlib.sha256):
    h_len = hash_func().digest_size
    if mask_len_bytes > (2**32 * h_len):
        raise ValueError("mask too long")
    t = b""
    num_blocks = math.ceil(mask_len_bytes / h_len)
    for counter in range(num_blocks):
        c_bytes = counter.to_bytes(4, 'big')
        t += hash_func(seed_bytes + c_bytes).digest()
    return t[:mask_len_bytes]

def bytes_xor(b1, b2):
    if len(b1) != len(b2):
        raise ValueError("bytes_xor: 输入的字节串长度必须相同")
    return bytes(x ^ y for x, y in zip(b1, b2))


def generate_elgamal_parameters(key_size_bits=2048):
    """
    生成一般的 ElGamal 参数 (p, g)
    """
    print(f"为 ElGamal 生成 {key_size_bits} 位素数 p 和生成元 g...")

    p = generate_prime_number(key_size_bits)
    print(f"已生成素数 p = {p}")

    while True:
        # 修改：使用 secrets 模块来安全地寻找生成元 g
        g = secrets.randbelow(p - 3) + 2 # 范围 [2, p-2]
        if power(g, (p - 1) // 2, p) != 1:
            print(f"已找到生成元 g = {g}")
            break

    return p, None, g